{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary modules\n",
    "from classy import Class\n",
    "from math import pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#####################################################\n",
    "#\n",
    "# Cosmological parameters and other CLASS parameters\n",
    "#\n",
    "#####################################################\n",
    "common_settings = {# LambdaCDM parameters\n",
    "                   'h':0.67810,\n",
    "                   'omega_b':0.02238280,\n",
    "                   'omega_cdm':0.12038,\n",
    "                   'A_s':2.100549e-09,\n",
    "                   'n_s': 0.9660499,\n",
    "                   'tau_reio':0.05430842,\n",
    "                   # output and precision parameters\n",
    "                   'output':'tCl,mTk,vTk',\n",
    "                   'l_max_scalars':5000,\n",
    "                   'P_k_max_1/Mpc':10.0,\n",
    "                   'gauge':'newtonian'\n",
    "                   }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "###############\n",
    "#\n",
    "# call CLASS a first time just to compute z_rec (will compute transfer functions at default: z=0)\n",
    "#\n",
    "###############\n",
    "M = Class()\n",
    "M.set(common_settings)\n",
    "M.compute()\n",
    "derived = M.get_current_derived_parameters(['z_rec','tau_rec','conformal_age'])\n",
    "print (derived.keys())\n",
    "z_rec = derived['z_rec']\n",
    "z_rec = int(1000.*z_rec)/1000. # round down at 4 digits after coma\n",
    "print ('z_rec=',z_rec)\n",
    "#\n",
    "# In the last figure the x-axis will show l/(tau_0-tau_rec), so we need (tau_0-tau_rec) in units of [Mpc/h]\n",
    "#\n",
    "tau_0_minus_tau_rec_hMpc = (derived['conformal_age']-derived['tau_rec'])*M.h()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "################\n",
    "#\n",
    "# call CLASS again for the perturbations (will compute transfer functions at input value z_rec)\n",
    "#\n",
    "################\n",
    "M.empty() # reset input parameters to default, before passing a new parameter set\n",
    "M.set(common_settings)\n",
    "M.set({'z_pk':z_rec})\n",
    "M.compute()\n",
    "#\n",
    "# save the total Cl's (we will plot them in the last step)\n",
    "#\n",
    "cl_tot = M.raw_cl(5000)\n",
    "#\n",
    "#\n",
    "# load transfer functions at recombination\n",
    "#\n",
    "one_time = M.get_transfer(z_rec)\n",
    "print (one_time.keys())\n",
    "k = one_time['k (h/Mpc)']\n",
    "Theta0 = 0.25*one_time['d_g']\n",
    "phi = one_time['phi']\n",
    "psi = one_time['psi']\n",
    "theta_b = one_time['t_b']\n",
    "# compute related quantitites\n",
    "R = 3./4.*M.Omega_b()/M.Omega_g()/(1+z_rec)  # R = 3/4 * (rho_b/rho_gamma) at z_rec\n",
    "zero_point = -(1.+R)*psi                     # zero point of oscillations: -(1.+R)*psi\n",
    "Theta0_amp = max(Theta0.max(),-Theta0.min()) # Theta0 oscillation amplitude (for vertical scale of plot)\n",
    "print ('At z_rec: R=',R,', Theta0_amp=',Theta0_amp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use table of background quantitites to find the wavenumbers corresponding to\n",
    "# Hubble crossing (k = 2 pi a H), sound horizon crossing (k = 2pi / rs)\n",
    "#\n",
    "background = M.get_background() # load background table\n",
    "print (background.keys())\n",
    "#\n",
    "background_tau = background['conf. time [Mpc]'] # read confromal times in background table\n",
    "background_z = background['z'] # read redshift\n",
    "background_kh = 2.*pi*background['H [1/Mpc]']/(1.+background['z'])/M.h() # read kh = 2pi aH = 2pi H/(1+z) converted to [h/Mpc]\n",
    "background_ks = 2.*pi/background['comov.snd.hrz.']/M.h() # read ks = 2pi/rs converted to [h/Mpc]\n",
    "#\n",
    "# define interpolation functions; we want the value of tau when the argument is equal to 2pi\n",
    "#\n",
    "from scipy.interpolate import interp1d\n",
    "kh_at_tau = interp1d(background_tau,background_kh)\n",
    "ks_at_tau = interp1d(background_tau,background_ks)\n",
    "#\n",
    "# finally get these scales\n",
    "#\n",
    "tau_rec = derived['tau_rec']\n",
    "kh = kh_at_tau(tau_rec)\n",
    "ks = ks_at_tau(tau_rec)\n",
    "print ('at tau_rec=',tau_rec,', kh=',kh,', ks=',ks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#####################\n",
    "#\n",
    "# call CLASS with TSW (intrinsic temperature + Sachs-Wolfe) and save\n",
    "#\n",
    "#####################\n",
    "M.empty()           # clean input\n",
    "M.set(common_settings) # new input\n",
    "M.set({'temperature contributions':'tsw'})\n",
    "M.compute()\n",
    "cl_TSW = M.raw_cl(5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######################\n",
    "#\n",
    "# call CLASS with early ISW and save\n",
    "#\n",
    "######################\n",
    "M.empty()\n",
    "M.set(common_settings)\n",
    "M.set({'temperature contributions':'eisw'})\n",
    "M.compute()\n",
    "cl_eISW = M.raw_cl(5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######################\n",
    "#\n",
    "# call CLASS with late ISW and save\n",
    "#\n",
    "######################\n",
    "M.empty()\n",
    "M.set(common_settings)\n",
    "M.set({'temperature contributions':'lisw'})\n",
    "M.compute()\n",
    "cl_lISW = M.raw_cl(5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######################\n",
    "#\n",
    "# call CLASS with Doppler and save\n",
    "#\n",
    "######################\n",
    "M.empty()\n",
    "M.set(common_settings)\n",
    "M.set({'temperature contributions':'dop'})\n",
    "M.compute()\n",
    "cl_Doppler = M.raw_cl(5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# modules and esthetic definitions for the plots\n",
    "#\n",
    "# uncomment to get plots displayed in notebook\n",
    "%matplotlib inline\n",
    "#\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "#\n",
    "font = {'size'   : 16, 'family':'STIXGeneral'}\n",
    "axislabelfontsize='large'\n",
    "matplotlib.rc('font', **font)\n",
    "matplotlib.mathtext.rcParams['legend.fontsize']='medium'\n",
    "plt.rcParams[\"figure.figsize\"] = [8.0,6.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#################\n",
    "#\n",
    "# start plotting\n",
    "#\n",
    "#################\n",
    "#\n",
    "fig, (ax_Tk, ax_Tk2, ax_Cl) = plt.subplots(3,sharex=True,figsize=(8,12))\n",
    "fig.subplots_adjust(hspace=0)\n",
    "##################\n",
    "#\n",
    "# first figure with transfer functions\n",
    "#\n",
    "##################\n",
    "ax_Tk.set_xlim([3.e-4,0.5])\n",
    "ax_Tk.set_ylim([-1.1*Theta0_amp,1.1*Theta0_amp])\n",
    "ax_Tk.tick_params(axis='x',which='both',bottom='off',top='on',labelbottom='off',labeltop='on')\n",
    "ax_Tk.set_xlabel(r'$\\mathrm{k} \\,\\,\\,  \\mathrm{[h/Mpc]}$')\n",
    "ax_Tk.set_ylabel(r'$\\mathrm{Transfer}(\\tau_\\mathrm{dec},k)$')\n",
    "ax_Tk.xaxis.set_label_position('top')\n",
    "ax_Tk.grid()\n",
    "#\n",
    "ax_Tk.axvline(x=kh,color='r')\n",
    "ax_Tk.axvline(x=ks,color='y')\n",
    "#\n",
    "ax_Tk.annotate(r'Hubble cross.',\n",
    "                xy=(kh,0.8*Theta0_amp),\n",
    "                xytext=(0.15*kh,0.9*Theta0_amp),\n",
    "                arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "ax_Tk.annotate(r'sound hor. cross.',\n",
    "                 xy=(ks,0.8*Theta0_amp),\n",
    "                 xytext=(1.3*ks,0.9*Theta0_amp),\n",
    "                 arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "#\n",
    "ax_Tk.semilogx(k,psi,'y-',label=r'$\\psi$')\n",
    "ax_Tk.semilogx(k,phi,'r-',label=r'$\\phi$')\n",
    "ax_Tk.semilogx(k,zero_point,'k:',label=r'$-(1+R)\\psi$')\n",
    "ax_Tk.semilogx(k,Theta0,'b-',label=r'$\\Theta_0$')\n",
    "ax_Tk.semilogx(k,(Theta0+psi),'c',label=r'$\\Theta_0+\\psi$')\n",
    "ax_Tk.semilogx(k,theta_b,'g-',label=r'$\\theta_b$')\n",
    "#\n",
    "ax_Tk.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "#######################\n",
    "#\n",
    "# second figure with transfer functions squared\n",
    "#\n",
    "#######################\n",
    "ax_Tk2.set_xlim([3.e-4,0.5])\n",
    "ax_Tk2.tick_params(axis='x',which='both',bottom='off',top='off',labelbottom='off',labeltop='off')\n",
    "ax_Tk2.set_ylabel(r'$\\mathrm{Transfer}(\\tau_\\mathrm{dec},k)^2$')\n",
    "ax_Tk2.grid()\n",
    "#\n",
    "ax_Tk2.semilogx(k,(Theta0+psi)**2,'c',label=r'$(\\Theta_0+\\psi)^2$')\n",
    "#\n",
    "ax_Tk2.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "########################\n",
    "#\n",
    "# third figure with all contributions to Cls\n",
    "#\n",
    "# We already computed each contribution (TSW, earlyISW, lateISW, Doppler, total)\n",
    "# Note that there is another contribution from polarisation. We don't plot it because it is\n",
    "# too small to be seen, however it is included by default in the total.\n",
    "#\n",
    "# After each step we will save the figure (to get intermediate figures that can be used in slides)\n",
    "#\n",
    "#########################\n",
    "# presentation settings\n",
    "ax_Cl.set_xlim([3.e-4,0.5])\n",
    "ax_Cl.set_ylim([0.,8.])\n",
    "ax_Cl.set_xlabel(r'$\\ell/(\\tau_0-\\tau_{rec}) \\,\\,\\, \\mathrm{[h/Mpc]}$')\n",
    "ax_Cl.set_ylabel(r'$\\ell (\\ell+1) C_l^{TT} / 2 \\pi \\,\\,\\, [\\times 10^{10}]$')\n",
    "ax_Cl.tick_params(axis='x',which='both',bottom='on',top='off',labelbottom='on',labeltop='off')\n",
    "ax_Cl.grid()\n",
    "#\n",
    "# plot and save with TSW\n",
    "#\n",
    "ax_Cl.semilogx(cl_TSW['ell']/tau_0_minus_tau_rec_hMpc,1.e10*cl_TSW['ell']*(cl_TSW['ell']+1.)*cl_TSW['tt']/2./pi,'c-',label=r'$\\mathrm{T+SW}$')\n",
    "#\n",
    "ax_Cl.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "fig.savefig('one_time_with_cl_1.pdf',bbox_inches='tight')\n",
    "#\n",
    "# plot and save with additionally early ISW and late ISW\n",
    "#\n",
    "ax_Cl.semilogx(cl_eISW['ell']/tau_0_minus_tau_rec_hMpc,1.e10*cl_eISW['ell']*(cl_eISW['ell']+1.)*cl_eISW['tt']/2./pi,'r-',label=r'$\\mathrm{early} \\,\\, \\mathrm{ISW}$')\n",
    "ax_Cl.semilogx(cl_lISW['ell']/tau_0_minus_tau_rec_hMpc,1.e10*cl_lISW['ell']*(cl_lISW['ell']+1.)*cl_lISW['tt']/2./pi,'y-',label=r'$\\mathrm{late} \\,\\, \\mathrm{ISW}$')\n",
    "#\n",
    "ax_Cl.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "fig.savefig('one_time_with_cl_2.pdf',bbox_inches='tight')\n",
    "#\n",
    "# plot and save with additionally Doppler\n",
    "#\n",
    "ax_Cl.semilogx(cl_Doppler['ell']/tau_0_minus_tau_rec_hMpc,1.e10*cl_Doppler['ell']*(cl_Doppler['ell']+1.)*cl_Doppler['tt']/2./pi,'g-',label=r'$\\mathrm{Doppler}$')\n",
    "#\n",
    "ax_Cl.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "fig.savefig('one_time_with_cl_3.pdf',bbox_inches='tight')\n",
    "#\n",
    "# plot and save with additionally total Cls\n",
    "#\n",
    "ax_Cl.semilogx(cl_tot['ell']/tau_0_minus_tau_rec_hMpc,1.e10*cl_tot['ell']*(cl_tot['ell']+1.)*cl_tot['tt']/2./pi,'k-',label=r'$\\mathrm{Total}$')\n",
    "#\n",
    "ax_Cl.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "fig.savefig('one_time_with_cl_tot.pdf',bbox_inches='tight')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
